import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";

import { readFileSync } from 'fs';
import path from 'path';

import { McpToolCallResponse } from './mcpType.js'
import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'

class MCPServerInfo {
    public name: string = "";
    public tools: Map<string, MCPToolInfo> = new Map();
    public mcp?: Client;
    // 启动配置
    public command: string = "";
    public args?: Array<string>;
}

class MCPToolInfo {
    public name?: string;
    public description?: string;
    public input_schema?: any;
}

export interface MCPConfig {
    configFilePath: string;
}

export class MCPClient {
    public mcpConfig: MCPConfig;
    public servers: Map<string, MCPServerInfo> = new Map();

    public static async create(mcpConfig: MCPConfig):Promise<MCPClient> {
        const mcpClient = new MCPClient(mcpConfig);
        await mcpClient.parseFile();
        return mcpClient;
    }

    constructor(mcpConfig: MCPConfig) {
        this.mcpConfig = mcpConfig;
    }

    private async parseFile(): Promise<void> {
        const fullPath = path.resolve(this.mcpConfig.configFilePath);
        const mcpServerConfigStr = readFileSync(fullPath, 'utf-8');
        const mcpServerConfigJson = JSON.parse(mcpServerConfigStr);
        const mcpServerMap = mcpServerConfigJson.mcpServers;

        for (const [name, config] of Object.entries(mcpServerMap)) {
            if (config != null) {
                const mcpServerInfo = new MCPServerInfo();
                mcpServerInfo.name = name;
                const dataConfig = config as any;
                mcpServerInfo.command = dataConfig.command;
                mcpServerInfo.args = dataConfig.args;
                await this.connectToServer(mcpServerInfo);
            }
        }
    }

    private async connectToServer(mcpServerInfo: MCPServerInfo): Promise<void> {
        try {
            const mcp = new Client({ name: "mcp-client-cli", version: "1.0.0" });
            const transport: StdioClientTransport = new StdioClientTransport({
                command: mcpServerInfo.command,
                args: mcpServerInfo.args
            });
            mcp.connect(transport);
            const toolsResult = await mcp.listTools();

            mcpServerInfo.mcp = mcp;
            mcpServerInfo.tools = new Map<string, MCPToolInfo>();

            for (const tool of toolsResult.tools) {
                const mcpToolInfo = new MCPToolInfo();
                mcpToolInfo.name = tool.name;
                mcpToolInfo.description = tool.description;
                mcpToolInfo.input_schema = tool.inputSchema;
                mcpServerInfo.tools.set(tool.name, mcpToolInfo);
            }
            if (this.servers) {
                this.servers.set(mcpServerInfo.name, mcpServerInfo);
            }
        } catch (e) {
            console.log("Failed to connect to MCP server: ", e);
            throw e;
        }
    }

    public async processQuery(serverName: string, toolName: string, query: {}) {
        const mcpServerInfo = this.servers?.get(serverName);
        if (!mcpServerInfo || !mcpServerInfo.mcp) {
            throw new Error("mcp not found");
        }
        const mcpToolResult = await mcpServerInfo.mcp.callTool({
            name: toolName,
            arguments: query
        })

        if (Array.isArray(mcpToolResult.content)) {
            for (const item of mcpToolResult.content) {
                return item.text;
            }
        }
    }

    public async callTool(serverName: string, toolName: string, toolArguments?: Record<string, unknown>): Promise<McpToolCallResponse> {
        const connection = this.servers?.get(serverName);
        if (!connection || !connection.mcp) {
			throw new Error(
				`No connection found for server: ${serverName}. Please make sure to use MCP servers available under 'Connected MCP Servers'.`,
			)
		}

		let timeout: number = 60000;
		const result = await connection.mcp.request(
			{
				method: "tools/call",
				params: {
					name: toolName,
					arguments: toolArguments,
				},
			},
			CallToolResultSchema,
			{
				timeout,
			},
		)
		
		// Filter out unsupported content types
		const filteredContent = result.content.filter(item => 
			item.type === 'text' || item.type === 'image' || item.type === 'resource'
		)
		
		return {
			...result,
			content: filteredContent
		}
	}

    public async cleanup(): Promise<void> {
        if (this.servers) {
            for (const [serverName, serverInfo] of this.servers) {
                await serverInfo.mcp?.close();
            }
        }
    }
}
